import gradio as gr
from joblib import load
from PIL import Image
import numpy as np

# 加载模型
model = load("LGBMClassifier_model.pkl")

def create_features(image):
    """
    将 PIL 图像转换为模型需要的输入特征。
    """
    # 假设模型需要 224x224 的 RGB 图像
    image = image.resize((224, 224))
    image_array = np.array(image)
    
    # 归一化像素值到 [0, 1]
    image_array = image_array / 255.0
    
    # 将图像数组转换为一维数组，作为特征
    features = image_array.flatten()
    
    # 如果特征数量超过1024，我们可以使用PCA或其他方法进行降维
    # 如果特征数量少于1024，我们可以填充0或者使用其他方法增加特征数量
    if features.shape[0] > 1024:
        # 这里只是一个示例，实际上应该使用PCA或其他降维技术
        features = features[:1024]
    elif features.shape[0] < 1024:
        # 填充0以确保特征数量为1024
        features = np.pad(features, (0, 1024 - features.shape[0]), 'constant')
    
    return features

def predict_image(image):
    try:
        # 将图片转换为模型需要的输入特征
        features = create_features(image)
        
        # 确保特征数量与模型训练时一致
        assert features.shape[0] == 1024, "特征数量不匹配"
        
        # 使用模型进行预测
        prediction = model.predict([features])
        return "猫" if prediction[0] == 1 else "狗"
    except Exception as e:
        # 如果出现错误，返回错误消息
        return str(e)

# 创建 Gradio 接口
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),  # 接受 PIL 图片格式
    outputs="text",
    title="猫狗识别",
    description="上传一张图片，模型将预测它是猫还是狗。"
)

if __name__ == "__main__":
    demo.launch()